package club.babyf.nativeai.functioncalling.invoker;

import club.babyf.nativeai.entity.CustomHashMap;
import club.babyf.nativeai.functioncalling.exception.FunctionExecutionException;
import club.babyf.nativeai.functioncalling.functions.AbstractFunction;
import club.babyf.nativeai.functioncalling.registry.FunctionRegistry;
import club.babyf.nativeai.utils.JsonCustomUtil;
import cn.hutool.core.collection.CollUtil;
import com.fasterxml.jackson.core.type.TypeReference;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * FunctionClient.
 *
 * @author zhanglei.
 * @date 2025/2/21 11:18.
 * @description 函数客户端.
 */
@Component
@RequiredArgsConstructor
public class FunctionClient {

    private final FunctionRegistry registry;

    public boolean isInvokeFunction(String resBody) {
        // 获取第一个choice
        CustomHashMap<String, Object> choiceJson = getChoiceJson(resBody);
        // 获取message
        CustomHashMap<String, Object> message = choiceJson.getObject("message", CustomHashMap.class);
        // 获取tool_calls
        List<CustomHashMap<String, Object>> toolCalls = message.getList("tool_calls", new TypeReference<>() {
        });
        // 如果tool_calls为空，不需要处理functioncalling
        if (CollUtil.isEmpty(toolCalls)) {
            return false;
        }
        for (CustomHashMap<String, Object> toolCall : toolCalls) {
            // 获取function
            CustomHashMap<String, Object> function = toolCall.getObject("function", CustomHashMap.class);
            // 获取name
            String name = function.getString("name");
            // 判断是否存在该函数
            if (registry.hasFunction(name)) {
                return true;
            }
        }
        return false;
    }

    public CustomHashMap<String, Object> invokeFunction(String reqBody, String resBody) {
        // 获取第一个choice
        CustomHashMap<String, Object> choiceJson = getChoiceJson(resBody);
        // 获取message
        CustomHashMap<String, Object> message = choiceJson.getObject("message", CustomHashMap.class);
        // 获取tool_calls
        List<CustomHashMap<String, Object>> toolCalls = message.getList("tool_calls", new TypeReference<>() {
        });
        // 获取请求参数
        CustomHashMap<String, Object> reqBodyMap = JsonCustomUtil.parse(reqBody, CustomHashMap.class);
        // 获取messages
        List<CustomHashMap<String, Object>> messages = reqBodyMap.getList("messages", new TypeReference<>() {
        });
        // 将messageResult添加到messages中
        messages.add(message);
        // 遍历toolCalls
        for (CustomHashMap<String, Object> toolCall : toolCalls) {
            // 获取function
            CustomHashMap<String, Object> function = toolCall.getObject("function", CustomHashMap.class);
            // 获取name
            String name = function.getString("name");
            // 判断是否存在该函数
            if (!registry.hasFunction(name)) {
                continue;
            }
            // 获取arguments
            String arguments = function.getString("arguments");
            // 调用函数
            Object functionResult = doInvokeFunction(name, arguments);
            // 获取id
            String id = toolCall.getString("id");
            // 构建toolJson
            CustomHashMap<String, Object> toolJson = new CustomHashMap<>();
            toolJson.put("tool_call_id", id);
            toolJson.put("role", "tool");
            toolJson.put("name", name);
            toolJson.put("content", functionResult);
            messages.add(toolJson);
        }
        // 更新请求参数
        reqBodyMap.put("messages", messages);
        // 返回请求参数
        return reqBodyMap;
    }

    private CustomHashMap<String, Object> getChoiceJson(String resBody) {
        // 解析返回结果
        CustomHashMap<String, Object> resBodyMap = JsonCustomUtil.parse(resBody, CustomHashMap.class);
        // 获取choices
        List<CustomHashMap<String, Object>> choiceJsonArr = resBodyMap.getList("choices", new TypeReference<>() {
        });
        // 获取第一个choice
        return JsonCustomUtil.parse(JsonCustomUtil.toJsonString(choiceJsonArr.get(0)), CustomHashMap.class);
    }

    public Object doInvokeFunction(String functionName, String argsJson) throws FunctionExecutionException {
        AbstractFunction<Object, Object> function = (AbstractFunction<Object, Object>) registry.getFunction(functionName);
        if (function == null) {
            throw new FunctionExecutionException("FUNCTION_NOT_FOUND", "Function not registered: " + functionName);
        }
        Object params = function.parseParams(argsJson);
        return function.execute(params);
    }

}
